import os.path
import pickle
import random
import cv2
from tqdm.auto import tqdm
from matplotlib import pyplot as plt
from torch.utils.data import Dataset
from torchvision.transforms import transforms as T
import torch
from PIL import Image
from torchvision.datasets import MNIST
from multiprocessing.pool import ThreadPool
from ylcls.utils import LOGGER, colorstr
NUM_THREADS = min(8, max(1, os.cpu_count() - 1))
class CLSDataset(Dataset):
    def __init__(self, config, mode = "train", mosaic = True, transforms = None, max_nums=None, cache=False):
        super(CLSDataset, self).__init__()
        self.config = config
        self.image_size = config.image_size
        self.mode = mode
        self.prefix = colorstr(f'{mode}: ')
        self.cache = cache
        self.images_cache_name = ''
        if mode == 'train':
            self.images_files_list, self.labels = config.train_images_file_list,  config.train_labels_list
        elif mode == 'valid':
            self.images_files_list, self.labels = config.valid_images_file_list, config.valid_labels_list
        elif mode == 'test':
            self.images_files_list, self.labels = config.test_images_file_list, config.test_labels_list
        else:
            self.images_files_list, self.labels = [], []
        if (max_nums != None):
            self.images_files_list, self.labels = self.images_files_list[:max_nums], self.labels[:max_nums]
        if cache:
            self.images_cache_name = self.get_cache_names()
        self.ni = len(self.images_files_list)
        self.images_list = [None] * self.ni
        self.mosaic = mosaic
        if transforms == None:
            if mode == 'train':
                if mosaic == False:
                    self.transforms = T.Compose(
                                [
                                    T.Resize(config.image_size + config.image_size // 4),
                                    T.CenterCrop(config.image_size),
                                    T.RandomHorizontalFlip(0.5),
                                    T.RandomVerticalFlip(0.5),
                                    T.ToTensor(),
                                    T.Normalize([0.5], [0.5]),
                                    # T.RandomErasing(),
                                ])
                else:
                    self.transforms = T.Compose(
                        [
                            T.Resize(config.image_size),
                            T.CenterCrop(config.image_size),
                            T.RandomHorizontalFlip(0.5),
                            T.RandomVerticalFlip(0.5),
                            T.ToTensor(),
                            T.Normalize([0.5], [0.5]),
                            # T.RandomErasing(),
                        ])
            else:
                self.transforms = T.Compose(
                    [
                        T.Resize(config.image_size),
                        T.ToTensor(),
                        T.Normalize([0.5], [0.5]),
                        # T.RandomErasing(),
                    ])
        else:
            self.transforms = transforms
        # if mode == "train":
        #     if mosaic == False:
        #         self.image_size = config.image_size + config.image_size // 4
        self.get_images()

    def __len__(self):
        return len(self.images_files_list)
    def __getitem__(self, index):
        if self.mosaic:
            img = self.get_mosaic4_images(index)
        else:
            img = self.images_list[index]
        return self.transforms(img), torch.tensor(self.labels[index])
    def get_cache_names(self):
        images_cache_name = os.path.join(self.config.dataset_path, f'{self.config.image_size}_images_cache.pickle')
        return images_cache_name

    def get_images(self):
        if not os.path.exists(self.images_cache_name):
            gb = 0
            fcn = cv2.imread
            results = ThreadPool(NUM_THREADS).imap(fcn, self.images_files_list)
            pbar = tqdm(enumerate(results), total=self.ni, mininterval=0.3)
            for i, x in pbar:
                x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)
                # x = cv2.resize(x, (self.image_size, self.image_size), interpolation=cv2.INTER_LINEAR)
                gb += x.nbytes
                x = Image.fromarray(x)
                self.images_list[i] = x  # im, hw_orig, hw_resized = load_image(self, i)
                pbar.desc = f"{self.mode} Caching images ({gb / 1E9:.1f}GB)"
            pbar.close()
            if self.cache:
                with open(self.images_cache_name, 'wb') as f:
                    pickle.dump(self.images_list, f)
                LOGGER.info(f'{self.prefix}images cache has been saved as pickle file!')
        else:
            with open(self.images_cache_name, 'rb') as f:
                self.images_list = pickle.load(f)
            LOGGER.info(f'{self.prefix}images cache has been loaded!')

    def get_mosaic4_images(self, index):
        w, h = self.image_size, self.image_size
        target = self.labels[index]
        labels = []
        for i, label in enumerate(self.labels):
            if label == target:
                labels.append(i)
        indices = random.choices(labels, k=4)  # 3 additional image indices
        if index not in indices:
            i = random.randint(0, 4 - 1)
            indices[i] = index
        random.shuffle(indices)
        grid = Image.new('RGB', size=(2 * w, 2 * h))
        images = [self.images_list[idx].resize((w, h), Image.ANTIALIAS) for idx in indices]
        for i, image in enumerate(images):
            grid.paste(image, box=(i % 2 * w, i // 2 * h))
        return grid
    def show_image(self, idx):
        img = Image.open(self.images_files_list[idx]).convert("RGB")
        plt.imshow(img)
        plt.axis("off")
        plt.show()

class MNISTCLSDataset(Dataset):
    def __init__(self, config, mode = "train", transforms = None, max_nums=None):
        super(MNISTCLSDataset, self).__init__()
        if transforms == None:
            self.transforms = T.Compose(
                [
                    T.Resize((config.image_size, config.image_size)),
                    T.RandomHorizontalFlip(0.5),
                    T.ToTensor(),
                    T.Normalize([0.5], [0.5]),
                ]
            )
        else:
            self.transforms = transforms
        self.images_labels = MNIST("..", True if mode=="train" else False, self.transforms, download=True)
        if(max_nums!=None):
            self.images_labels = self.images_labels[:max_nums]
    def __len__(self):
        return len(self.images_labels)
    def __getitem__(self, idx):
        return self.images_labels[idx]
    def show_image(self, idx):
        plt.imshow(self.images_labels[idx][0].permute(1,2,0).detach().cpu().numpy())
        plt.axis("off")
        plt.show()
if __name__ == '__main__':
    from config import garbage_clsconfig_dict, CLSConfig, get_config
    config = get_config(garbage_clsconfig_dict, CLSConfig)
    dataset = config.train_dataset
    pic = dataset.get_mosaic4_images(1)
    plt.imshow(pic)
    plt.axis("off")
    plt.show()
    print(pic.size)